Skip to content

Ollama 如何加入鉴权

背景介绍

Ollama 本身默认不带鉴权机制,任何能够访问 Ollama 服务端口的人都可以直接调用其 API。在某些场景下,这可能会带来安全风险,例如:

  • 将 Ollama 服务暴露在公网上时,容易被恶意访问和滥用。
  • 在团队内部共享 Ollama 服务时,希望控制访问权限,避免资源滥用。

为了解决这个问题,我们可以引入一个简单的代理服务器,在 Ollama 前面加一层鉴权逻辑。

实现原理

我们将使用 Python 的 Flask 框架创建一个轻量级的代理服务器。这个代理服务器的工作流程如下:

  1. 接收请求: 客户端(例如应用程序、curl 命令)将请求发送到代理服务器的地址和端口。
  2. 校验 Token: 代理服务器检查请求头中的 Authorization 字段,验证 Bearer Token 是否与预设的 Token 一致。
    • 如果 Token 无效或缺失,代理服务器直接返回 401 Unauthorized 错误。
    • 如果 Token 有效,进入下一步。
  3. 转发请求: 代理服务器将通过验证的请求(包括请求方法、路径、查询参数、请求头和请求体)转发给实际运行的 Ollama 服务。
  4. 返回响应: 代理服务器接收 Ollama 返回的响应,并将其原样(或经过适当处理)返回给客户端。

这种方式相当于在客户端和 Ollama 服务之间增加了一个"门卫",只有持有正确"钥匙"(Token)的请求才能通过。

代码示例

以下是一个使用 Flask 实现的简单 Ollama 鉴权代理脚本 (ollama_auth_proxy.py):

python
from flask import Flask, request, Response, jsonify, stream_with_context
import requests

app = Flask(__name__)

# 预期的 token 和目标服务器地址
EXPECTED_TOKEN = "YOUR_SECRET_TOKEN"  # 请替换成你的实际 token
TARGET_SERVER = "http://localhost:11434"  # Ollama 服务地址

# 捕获所有 URL 路径的请求,支持 GET、POST、PUT、DELETE、PATCH、OPTIONS 等方法
@app.route('/', defaults={'path': ''}, methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS'])
@app.route('/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS'])
def proxy(path):
    # === 1. 校验 Token ===
    auth_header = request.headers.get('Authorization')
    if not auth_header:
        return jsonify({"error": "Authorization header missing"}), 401

    parts = auth_header.split()
    if len(parts) != 2 or parts[0].lower() != "bearer" or parts[1] != EXPECTED_TOKEN:
        return jsonify({"error": "Invalid token"}), 401

    # === 2. 构造目标 URL 与请求参数 ===
    # 如请求 /chat/message,则转发到 http://localhost:11434/chat/message
    target_url = f"{TARGET_SERVER}/{path}"
    params = request.args.to_dict()

    # 构造请求头,过滤掉 Host、Authorization 以免干扰后端服务器
    headers = {
        key: value
        for key, value in request.headers
        if key.lower() not in ["host", "authorization"]
    }

    # 获取原始请求体(适用于 POST、PUT 等方法)
    data = request.get_data()

    # === 3. 发送请求到目标服务器,并开启流式响应 ===
    try:
        resp = requests.request(
            method=request.method,
            url=target_url,
            params=params,
            headers=headers,
            data=data,
            stream=True  # 开启流式响应,逐块获取数据
        )
    except requests.exceptions.RequestException as e:
        return jsonify({"error": "Error connecting to target server", "details": str(e)}), 502

    # === 4. 过滤目标响应的部分头信息 ===
    excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
    response_headers = [(name, value) for name, value in resp.headers.items() if name.lower() not in excluded_headers]

    # === 5. 定义生成器函数,实现对响应流的逐块传输 ===
    def generate():
        try:
            for chunk in resp.iter_content(chunk_size=8192):
                if chunk:
                    yield chunk
        except requests.exceptions.RequestException as e:
            # 注意:异常信息也转换为字节流返回
            yield f"Error streaming response: {str(e)}".encode("utf-8")

    # 使用 stream_with_context 保证在请求上下文中正确流式传输数据
    return Response(stream_with_context(generate()), status=resp.status_code, headers=response_headers)

if __name__ == '__main__':
    # 监听所有网卡,端口设置为5000(可根据需要修改)
    app.run(host='0.0.0.0', port=5000)

使用方法

1. 环境准备与依赖安装

你需要安装 Python 和 pip。推荐使用虚拟环境来隔离项目依赖。

选项一:使用 Conda 创建环境 (推荐)

bash
# 创建一个新的 conda 环境(例如命名为 ollama_proxy_env),指定 Python 版本
conda create --name ollama_proxy_env python=3.10 -y

# 激活环境
conda activate ollama_proxy_env

# 安装依赖
pip install Flask requests waitress

选项二:直接使用 pip (需要确保系统 Python 环境干净或使用 venv)

bash
# 创建虚拟环境 (可选但推荐)
# python -m venv ollama_proxy_venv
# source ollama_proxy_venv/bin/activate # Linux/macOS
# .\ollama_proxy_venv\Scripts\activate # Windows

# 安装依赖
pip install Flask requests waitress

2. 配置代理脚本

  • 修改 Token: 打开 ollama_auth_proxy.py 文件,将 EXPECTED_TOKEN = "YOUR_SECRET_TOKEN" 中的 "YOUR_SECRET_TOKEN" 替换为你自己设定的强密码或随机生成的 Token。请务必使用一个安全的 Token!
  • 确认 Ollama 地址: 检查 TARGET_SERVER = "http://localhost:11434" 是否是你实际运行 Ollama 服务的地址和端口。如果 Ollama 运行在不同的机器或端口上,请相应修改。
  • 确认代理端口: PROXY_PORT = 5000 是代理服务器监听的端口,确保这个端口没有被其他程序占用。

3. 运行代理脚本

在激活了安装好依赖的 Python 环境的终端中,运行脚本:

bash
python ollama_auth_proxy.py

如果一切正常,你会看到类似以下的输出,表示代理服务器已启动:

Starting Ollama Auth Proxy on http://0.0.0.0:5000
Forwarding requests to: http://localhost:11434
Expected Token: ********OKEN
INFO:waitress:Serving on http://0.0.0.0:5000

4. 配置客户端

现在,你需要将原来直接访问 Ollama (例如 http://localhost:11434) 的客户端,改为访问代理服务器的地址和端口 (例如 http://localhost:5000),并且在请求头中加入 Authorization: Bearer YOUR_SECRET_TOKEN

例如,如果你使用 OLLAMA_HOST 环境变量来配置 Ollama 客户端:

  • 修改 OLLAMA_HOST: 将其指向代理服务器,例如 export OLLAMA_HOST=http://localhost:5000
  • 添加 Header: 你的客户端代码需要添加 Authorization 请求头。具体实现方式取决于你使用的库或工具。

测试示例 (使用 curl)

假设你的代理服务器运行在 localhost:5000,Ollama 服务运行在 localhost:11434,设置的 Token 是 MySuperSecret123

1. 不带 Token 的请求 (预期失败):

bash
curl http://localhost:5000/api/tags

预期会收到 401 Unauthorized 错误和类似 { "error": "Authorization header missing" } 的 JSON 响应。

2. 带错误 Token 的请求 (预期失败):

bash
curl http://localhost:5000/api/tags -H "Authorization: Bearer WrongToken"

预期会收到 401 Unauthorized 错误和类似 { "error": "Invalid token" } 的 JSON 响应。

3. 带正确 Token 的请求 (预期成功):

bash
curl http://localhost:5000/api/tags -H "Authorization: Bearer MySuperSecret123"

预期会收到 Ollama 返回的正常模型列表,例如:

json
{
  "models": [
    {
      "name": "llama3:latest",
      "modified_at": "2024-04-24T10:35:00.123456Z",
      "size": 4108448771,
      "digest": "sha256:...",
      "details": { ... }
    }
    // ... 其他模型
  ]
}

4. 测试 POST 请求 (例如聊天):

bash
curl http://localhost:5000/api/chat \
  -H "Authorization: Bearer MySuperSecret123" \
  -H "Content-Type: application/json" \
  -d '{
    "model": "llama3",
    "messages": [
      { "role": "user", "content": "Why is the sky blue?" }
    ],
    "stream": false
  }'

预期会收到来自 Ollama 的聊天响应。

注意事项

  • 安全性:
    • 使用强 Token: 不要使用容易猜到的 Token。考虑使用 UUID 或者密码生成器生成复杂的 Token。
    • HTTPS: 如果代理服务器和客户端之间的通信需要通过不安全的网络(如公网),强烈建议使用 HTTPS 来加密通信。你可以使用像 Nginx 或 Caddy 这样的反向代理来轻松实现 HTTPS。
    • 环境变量/配置管理: 不要将 Token 硬编码在代码中。使用环境变量 (os.getenv('OLLAMA_AUTH_TOKEN', 'DEFAULT_TOKEN')) 或配置文件来管理敏感信息。
  • 生产部署:
    • WSGI 服务器: 不要使用 Flask 内建的开发服务器 (app.run()) 进行生产部署。它性能不高且不够稳定。推荐使用 waitress (简单,跨平台) 或 gunicorn (功能更强,主要用于 Linux/macOS)。
    • 进程管理: 使用进程管理工具(如 systemd, supervisor)来确保代理服务在后台稳定运行,并在意外退出时自动重启。
  • 性能: 这个简单的代理实现对于大多数个人或小团队场景足够。但对于高并发场景,可能需要考虑更专业的 API 网关解决方案或优化代理逻辑(例如使用异步框架如 FastAPI + uvicorn)。
  • 日志: 代码中加入了基本的日志记录。在生产环境中,建议配置更详细的日志记录和轮转,以便于问题排查。
  • 错误处理: 代码中包含了基本的错误处理(连接错误、超时)。可以根据需要进一步细化错误处理逻辑。

后台运行

为了方便管理,让 Ollama 服务和鉴权代理服务在后台稳定运行并记录日志,我们可以创建一个简单的启动脚本。

创建启动脚本

创建一个名为 start_ollama_services.sh 的脚本文件:

bash
#!/bin/bash

# 设置日志目录
LOG_DIR="$HOME/logs"
mkdir -p $LOG_DIR

# 设置时间戳(用于日志文件名)
TIMESTAMP=$(date +"%Y%m%d_%H%M%S")

# Ollama 服务配置
OLLAMA_MODELS_DIR="$HOME/autodl-tmp/models"
OLLAMA_LOG="$LOG_DIR/ollama_$TIMESTAMP.log"

# 代理服务配置
PROXY_SCRIPT="$HOME/projects/main.py"
PROXY_LOG="$LOG_DIR/ollama_proxy_$TIMESTAMP.log"

# 启动 Ollama 服务
echo "启动 Ollama 服务..."
OLLAMA_MODELS=$OLLAMA_MODELS_DIR nohup ollama serve > $OLLAMA_LOG 2>&1 &
OLLAMA_PID=$!
echo "Ollama 服务已启动,PID: $OLLAMA_PID,日志: $OLLAMA_LOG"

# 等待 Ollama 服务完全启动
echo "等待 Ollama 服务初始化 (5秒)..."
sleep 5

# 启动鉴权代理服务
echo "启动鉴权代理服务..."
cd $(dirname $PROXY_SCRIPT)
nohup python $(basename $PROXY_SCRIPT) > $PROXY_LOG 2>&1 &
PROXY_PID=$!
echo "鉴权代理服务已启动,PID: $PROXY_PID,日志: $PROXY_LOG"

# 保存 PID 到文件中,方便后续停止服务
echo $OLLAMA_PID > $LOG_DIR/ollama.pid
echo $PROXY_PID > $LOG_DIR/ollama_proxy.pid

echo "服务启动完成!"
echo "Ollama 服务 PID: $OLLAMA_PID"
echo "鉴权代理 PID: $PROXY_PID"
echo "查看 Ollama 日志: tail -f $OLLAMA_LOG"
echo "查看代理服务日志: tail -f $PROXY_LOG"

创建停止脚本

同时,为了方便停止服务,可以创建一个 stop_ollama_services.sh 脚本:

bash
#!/bin/bash

# 设置日志目录
LOG_DIR="$HOME/logs"

# 检查 PID 文件
if [ -f "$LOG_DIR/ollama.pid" ]; then
    OLLAMA_PID=$(cat $LOG_DIR/ollama.pid)
    echo "停止 Ollama 服务 (PID: $OLLAMA_PID)..."
    kill -15 $OLLAMA_PID 2>/dev/null || echo "Ollama 服务已不在运行"
    rm -f $LOG_DIR/ollama.pid
else
    echo "找不到 Ollama PID 文件"
fi

if [ -f "$LOG_DIR/ollama_proxy.pid" ]; then
    PROXY_PID=$(cat $LOG_DIR/ollama_proxy.pid)
    echo "停止鉴权代理服务 (PID: $PROXY_PID)..."
    kill -15 $PROXY_PID 2>/dev/null || echo "鉴权代理服务已不在运行"
    rm -f $LOG_DIR/ollama_proxy.pid
else
    echo "找不到鉴权代理 PID 文件"
fi

echo "服务已停止"

使用方法

  1. 首先给脚本添加执行权限:
bash
chmod +x start_ollama_services.sh
chmod +x stop_ollama_services.sh
  1. 启动服务:
bash
./start_ollama_services.sh
  1. 查看服务状态和日志:
bash
# 查看 Ollama 服务日志
tail -f ~/logs/ollama_*.log

# 查看代理服务日志
tail -f ~/logs/ollama_proxy_*.log

# 检查进程是否在运行
ps -p $(cat ~/logs/ollama.pid) || echo "Ollama 服务未运行"
ps -p $(cat ~/logs/ollama_proxy.pid) || echo "鉴权代理服务未运行"
  1. 停止服务:
bash
./stop_ollama_services.sh

设置开机自启(可选)

如果你希望服务在系统启动时自动运行,可以考虑:

  1. 使用 systemd(Linux)

    创建 systemd 服务文件,例如 /etc/systemd/system/ollama-services.service

    ini
    [Unit]
    Description=Ollama and Auth Proxy Services
    After=network.target
    
    [Service]
    Type=forking
    User=你的用户名
    ExecStart=/path/to/start_ollama_services.sh
    ExecStop=/path/to/stop_ollama_services.sh
    Restart=on-failure
    
    [Install]
    WantedBy=multi-user.target

    然后启用并启动服务:

    bash
    sudo systemctl enable ollama-services.service
    sudo systemctl start ollama-services.service
  2. 使用 crontab

    bash
    crontab -e

    添加以下行:

    @reboot /path/to/start_ollama_services.sh

通过以上配置,你可以方便地在后台运行 Ollama 和鉴权代理服务,并轻松管理这些服务的启动、停止和日志记录。